package com.baijiaoxi.common.model.concurrent;

import com.baijiaoxi.common.model.local.ThreadLocalUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.*;

/**
 * 自定义线程池
 * 必须指定名称
 * 默认工厂类threadFactory，不支持外部自定义
 * 可以自定义拒绝策略
 */
@Slf4j
public class BaiThreadPoolExecutor extends ThreadPoolExecutor {


    private String poolName;

    public String getPoolName() {
        return poolName;
    }

    /**
     * 全部指定
     *
     * @param corePoolSize
     * @param maximumPoolSize
     * @param keepAliveTime
     * @param unit
     * @param workQueue
     * @param handler
     */
    public BaiThreadPoolExecutor(String poolName, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        this.poolName = poolName;
    }


    /**
     * Executes the given task sometime in the future.  The task
     * may execute in a new thread or in an existing pooled thread.
     * <p>
     * If the task cannot be submitted for execution, either because this
     * executor has been shutdown or because its capacity has been reached,
     * the task is handled by the current {@code RejectedExecutionHandler}.
     *
     * @param command the task to execute
     * @throws RejectedExecutionException at discretion of
     *                                    {@code RejectedExecutionHandler}, if the task
     *                                    cannot be accepted for execution
     * @throws NullPointerException       if {@code command} is null
     */
    @Override
    public void execute(Runnable command) {
        Object va = ThreadLocalUtil.get();
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        super.execute(new Runnable() {
            @Override
            public void run() {
                if (va != null)
                    ThreadLocalUtil.set(va);
                if (contextMap != null)
                    MDC.setContextMap(contextMap);
                try {
                    command.run();
                } catch (Exception e) {
                    log.error("线程执行错误", e);
                } finally {
                    ThreadLocalUtil.clear();
                    MDC.clear();
                }
            }
        });
    }

    /**
     * @param task
     * @throws RejectedExecutionException {@inheritDoc}
     * @throws NullPointerException       {@inheritDoc}
     */
    @Override
    public <T> Future<T> submit(Callable<T> task) {
        Object va = ThreadLocalUtil.get();
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        Future<T> future = super.submit(new Callable<T>() {
            /**
             * Computes a result, or throws an exception if unable to do so.
             *
             * @return computed result
             * @throws Exception if unable to compute a result
             */
            @Override
            public T call() throws Exception {
                if (va != null)
                    ThreadLocalUtil.set(va);
                if (contextMap != null)
                    MDC.setContextMap(contextMap);
                try {
                    return task.call();
                } catch (Exception e) {
                    log.error("线程执行错误", e);
                    return null;
                } finally {
                    ThreadLocalUtil.clear();
                    MDC.clear();
                }

            }
        });
        return future;
    }

}
